package com.example.im.message;

import com.example.im.entity.ChatMessage;
import com.example.im.protocol.Message;
import com.example.im.config.WebSocketHandler;
import com.example.im.service.ChatMessageService;
import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Lazy;
import org.springframework.stereotype.Component;

import javax.annotation.PreDestroy;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.*;

@Slf4j
@Component
public class MessageManager {
    // 待投递的消息任务Map，key为消息ID
    private final Map<String, DeliveryTask> pendingDeliveries = new ConcurrentHashMap<>();
    // 重试任务调度器
    private final ScheduledExecutorService scheduler = Executors.newScheduledThreadPool(1);
    
    // 最大重试次数
    private static final int MAX_RETRIES = 3;
    // 重试间隔（毫秒）
    private static final int RETRY_INTERVAL = 2000;

    @Autowired
    @Lazy
    private WebSocketHandler webSocketHandler;

    @Autowired
    private ChatMessageService chatMessageService;

    @Data
    @AllArgsConstructor
    private static class DeliveryTask {
        private Message message;
        private int retryCount;
        private ScheduledFuture<?> future;
    }

    /**
     * 处理新接收的消息
     * 1. 对于聊天消息：保存并发送确认，然后开始投递
     * 2. 对于确认消息：处理客户端确认
     */
    public void handleIncomingMessage(Message message) {
        try {
            if (message.getType() == Message.Type.CHAT) {
                // 1. 保存消息并发送确认
                saveMessageAndSendAck(message);
                // 2. 开始消息投递
                startMessageDelivery(message);
            } else if (message.getType() == Message.Type.CLIENT_ACK) {
                // 处理客户端确认消息
                handleClientAck(message);
            }
        } catch (Exception e) {
            log.error("Error handling incoming message", e);
            sendFailureAck(message);
        }
    }

    /**
     * 保存消息到数据库并发送服务器确认
     */
    private void saveMessageAndSendAck(Message message) {
        // 保存消息到数据库
        ChatMessage chatMessage = new ChatMessage();
        chatMessage.setMessageId(message.getMessageId());
        chatMessage.setFrom(message.getFrom());
        chatMessage.setTo(message.getTo());
        chatMessage.setContent(message.getContent());
        chatMessage.setTimestamp(message.getTimestamp());
        chatMessage.setStatus(ChatMessage.Status.SENT);
        chatMessageService.saveMessage(chatMessage);

        // 发送服务器确认
        sendServerAck(message.getMessageId(), message.getFrom(), Message.Status.SENT);
    }

    /**
     * 开始消息投递流程
     * 创建投递任务并进行首次投递尝试
     */
    private void startMessageDelivery(Message message) {
        DeliveryTask task = new DeliveryTask(message, 0, null);
        pendingDeliveries.put(message.getMessageId(), task);
        attemptDelivery(message.getMessageId());
    }

    /**
     * 尝试投递消息
     * 如果投递失败，会触发重试机制
     */
    private void attemptDelivery(String messageId) {
        DeliveryTask task = pendingDeliveries.get(messageId);
        if (task == null) return;

        try {
            webSocketHandler.sendMessage(task.getMessage());
        } catch (Exception e) {
            log.error("Failed to deliver message: {}", messageId, e);
            scheduleRetry(messageId);
        }
    }

    /**
     * 安排重试任务
     * 如果达到最大重试次数，会标记消息为失败
     */
    private void scheduleRetry(String messageId) {
        DeliveryTask task = pendingDeliveries.get(messageId);
        if (task == null) return;

        if (task.getRetryCount() >= MAX_RETRIES) {
            handleDeliveryFailed(messageId);
            return;
        }

        task.setRetryCount(task.getRetryCount() + 1);
        ScheduledFuture<?> future = scheduler.schedule(
            () -> attemptDelivery(messageId),
            RETRY_INTERVAL,
            TimeUnit.MILLISECONDS
        );
        task.setFuture(future);
    }

    /**
     * 处理投递失败的情况
     * 更新消息状态并通知发送方
     */
    private void handleDeliveryFailed(String messageId) {
        DeliveryTask task = pendingDeliveries.remove(messageId);
        if (task != null) {
            chatMessageService.updateMessageStatus(messageId, ChatMessage.Status.FAILED);
            sendServerAck(messageId, task.getMessage().getFrom(), Message.Status.FAILED);
        }
    }

    /**
     * 处理客户端确认消息
     * 1. 取消重试任务
     * 2. 更新消息状态
     * 3. 转发确认给原发送者
     */
    private void handleClientAck(Message ackMessage) {
        String messageId = ackMessage.getAckMessageId();
        DeliveryTask task = pendingDeliveries.remove(messageId);
        
        if (task != null) {
            // 取消重试任务
            if (task.getFuture() != null) {
                task.getFuture().cancel(false);
            }

            // 更新消息状态
            chatMessageService.updateMessageStatus(messageId, ChatMessage.Status.DELIVERED);

            // 转发送达确认给原发送者
            sendDeliveryAck(messageId, task.getMessage().getFrom(), ackMessage.getFrom());
        }
    }

    private void sendServerAck(String messageId, String to, Message.Status status) {
        Message ackMessage = new Message();
        ackMessage.setType(Message.Type.SERVER_ACK);
        ackMessage.setAckMessageId(messageId);
        ackMessage.setTo(to);
        ackMessage.setStatus(status);
        ackMessage.setTimestamp(System.currentTimeMillis());

        try {
            webSocketHandler.sendMessage(ackMessage);
        } catch (Exception e) {
            log.error("Failed to send server ACK", e);
        }
    }

    private void sendDeliveryAck(String messageId, String to, String from) {
        Message ackMessage = new Message();
        ackMessage.setType(Message.Type.DELIVERY_ACK);
        ackMessage.setAckMessageId(messageId);
        ackMessage.setTo(to);
        ackMessage.setFrom(from);
        ackMessage.setTimestamp(System.currentTimeMillis());

        try {
            webSocketHandler.sendMessage(ackMessage);
        } catch (Exception e) {
            log.error("Failed to send delivery ACK", e);
        }
    }

    private void sendFailureAck(Message message) {
        if (message.getMessageId() != null && message.getFrom() != null) {
            sendServerAck(message.getMessageId(), message.getFrom(), Message.Status.FAILED);
        }
    }

    @PreDestroy
    public void destroy() {
        scheduler.shutdown();
        try {
            if (!scheduler.awaitTermination(5, TimeUnit.SECONDS)) {
                scheduler.shutdownNow();
            }
        } catch (InterruptedException e) {
            scheduler.shutdownNow();
            Thread.currentThread().interrupt();
        }
    }
} 